import numpy as np
import matplotlib.pyplot as plt

fig,ax = plt.subplots()

data = np.loadtxt('SQETOTAL')

sqe = data[:,2].reshape((4096,20),order='F')
sqe[sqe<=0] = 1e-7


np.savetxt('E.txt',data[0:4096,0])
np.savetxt('0p2.txt',sqe[:,3])
np.savetxt('0p4.txt',sqe[:,7])
np.savetxt('0p8.txt',sqe[:,15])

E = data[0:4096,0]


kB = 8.61733e-2 #meV/K
occ = 1./(np.exp(E/(kB*600))-1) + 1

chi = sqe/occ[:,None]

#-1,2
vmin=-2.5
vmax=-0.5
fig = ax.imshow(np.log10(chi),origin='lower',vmin=vmin,vmax=vmax,aspect='auto',cmap='viridis',interpolation='none')


#xpos = np.linspace(-1,19,2)
ax.set_xlim(0,19)
#xticks = np.linspace(0,1,2,dtype='int')
xpos = [-1,3,7,11,15,19]
xticks = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
ax.set_xticks(xpos)
ax.set_xticklabels(xticks)
ax.set_xlabel('0 0 L (r.l.u.)',fontsize=20)

ypos = np.linspace(0,398,6)  # 398 20meV
yticks = np.linspace(0,20,6,dtype='int')
ax.set_yticks(ypos)
ax.set_yticklabels(yticks)
#ax.set_ylim(0,398)
ax.set_ylim(0,800)
ax.set_ylabel('Energy (meV)',fontsize=20)

ax.tick_params(which='both',labelsize=20)

cbar = plt.colorbar(fig,ticks=np.arange(vmin,vmax+0.1,0.5))
cbar.ax.tick_params(labelsize=20)
cbar.ax.set_yticklabels(np.arange(vmin,vmax+0.1,0.5))
plt.savefig('chi.png',format='png',bbox_inches='tight')
plt.show()
